/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2023-2023. All rights reserved.
 * Description: 集合通信算子AlltoAllVC,AllGather,ReduceScatter融合图优化头文件
 * Author: zhengjinwei
 * Create: 2023-1-29
 */

#ifndef HCOM_FUSION_OPTIMIZER_H
#define HCOM_FUSION_OPTIMIZER_H

#include <string>
#include "common/optimizer/graph_optimizer.h"
#include "common/optimizer/graph_optimizer_types.h"
#include "graph/compute_graph.h"
#include "hccl/hccl_types.h"
#include "base.h"

namespace hccl {
const string HCCL_FUSION_OPTIMIZER_NAME = "hccl_alltoallvc_fusion_optimizer";

class HcomFusionOptimizer : public ge::GraphOptimizer {
public:
    HcomFusionOptimizer();
    ~HcomFusionOptimizer() override;
    virtual ge::Status Initialize(const map<std::string, std::string>& options,
                                  ge::OptimizeUtility *const optimizeUtility) override;
    // close graphOptimizer
    ge::Status Finalize() override;
    // optimize original graph for FE quant optimize
    ge::Status OptimizeGraphPrepare(ge::ComputeGraph& graph) override;
    // optimize original graph, using in graph preparation stage
    ge::Status OptimizeOriginalGraph(ge::ComputeGraph& graph) override;
    // optimize fused graph
    ge::Status OptimizeFusedGraph(ge::ComputeGraph& graph) override;
    // optimize whole graph, using after graph merged stage
    ge::Status OptimizeWholeGraph(ge::ComputeGraph& graph) override;
    // get attribute of graph optimizer
    ge::Status GetAttributes(ge::GraphOptimizerAttribute& attrs) const override;
protected:
    HcclResult HcomOptimizeOriginalGraph(ge::ComputeGraph& graph);
private:
    HcclResult FuseHcomAlltoAllVCNode(ge::ComputeGraph& graph);
    HcclResult FuseHcomAllgatherNode(ge::ComputeGraph& graph);
    HcclResult FuseHcomReduceScatterNode(ge::ComputeGraph& graph);
    HcclResult HcomOptimizeSetAttr(ge::ComputeGraph& graph);
};
}
#endif